
#include <cuda_runtime.h>
#include <iostream>
#include <vector>
#include <random>

// 两级直方图设计已成功实现：
// 分 10 个粗略桶
// 每个桶用线程块独立统计局部直方图（共享内存）
// 合并到 100 类全局直方图
// 使用原子操作保证正确性

__global__
void twoLevelHisto(const unsigned int* const vals, unsigned int* const histo, int numVals) {
    int coarseBin = blockIdx.x;
    if (coarseBin >= 10) return;

    extern __shared__ unsigned int sharedHisto[];
    
    // 初始化共享内存
    for (int i = threadIdx.x; i < 10; i += blockDim.x) {
        sharedHisto[i] = 0;
    }
    __syncthreads();

    // 每个 block 扫描整个数组
    for (int i = threadIdx.x; i < numVals; i += blockDim.x) {
        unsigned int val = vals[i];
        int bucket = val / 10;  // 确定它属于哪个 coarseBin（十位数）
        int fineBin = val % 10; // 确定它在该 coarseBin 中的个位数位置（0~9）

        if (bucket == coarseBin) {
            atomicAdd(&sharedHisto[fineBin], 1);
        }
    }

    __syncthreads();

    // 写回全局直方图
    for (int i = threadIdx.x; i < 10; i += blockDim.x) {
        if (sharedHisto[i] > 0) {
            int globalBin = coarseBin * 10 + i;
            histo[globalBin] += sharedHisto[i];
            // atomicAdd(&histo[globalBin], sharedHisto[i]);
        }
    }
}

// CPU 参考实现：计算直方图（Golden Model）
void cpuHistogram(const std::vector<unsigned int>& vals, std::vector<unsigned int>& histo, int numBins) {
    histo.assign(numBins, 0);
    for (unsigned int val : vals) {
        if (val < numBins) {
            histo[val]++;
        }
        // 假设 val 在 [0,99] 合法范围内
    }
}

// 打印直方图（非零项）
void printHistogram(const std::vector<unsigned int>& histo, const std::string& name) {
    std::cout << "\n" << name << ":\n";
    bool hasPrinted = false;
    for (int i = 0; i < histo.size(); ++i) {
        if (histo[i] > 0) {
            std::cout << "  Bin[" << i << "] = " << histo[i] << "\n";
            hasPrinted = true;
        }
    }
    if (!hasPrinted) {
        std::cout << "  (all zeros)\n";
    }
}

int main() {
    const int numVals = 1000;
    const int numFineBins = 100;
    const int numCoarseBins = 10;
    const int threadsPerBlock = 256;

    // 1. 生成输入数据：1000 个 [0,99] 的随机整数
    std::vector<unsigned int> h_vals(numVals);
    std::random_device rd;
    std::mt19937 gen(rd());
    std::uniform_int_distribution<> dis(0, 99);

    for (int i = 0; i < numVals; i++) {
        h_vals[i] = dis(gen);
    }

    // 2. 分配设备内存
    unsigned int *d_vals, *d_histo;
    cudaMalloc(&d_vals, numVals * sizeof(unsigned int));
    cudaMalloc(&d_histo, numFineBins * sizeof(unsigned int));

    // 3. 拷贝输入数据到设备
    cudaMemcpy(d_vals, h_vals.data(), numVals * sizeof(unsigned int), cudaMemcpyHostToDevice);

    // 4. 初始化设备端直方图为 0
    cudaMemset(d_histo, 0, numFineBins * sizeof(unsigned int));

    // 创建 CUDA 事件用于计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    float time_gpu = 0.0f;
    

    // 5. 启动 kernel
    size_t sharedMemSize = numCoarseBins * sizeof(unsigned int);  // 每个 block 10 个 unsigned int
    int grid = numCoarseBins;
    cudaEventRecord(start);
    twoLevelHisto<<<grid, threadsPerBlock, sharedMemSize>>>(
        d_vals, d_histo, numVals);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&time_gpu, start, stop);

    // 6. 检查 kernel 是否出错
    cudaError_t err = cudaGetLastError();
    if (err != cudaSuccess) {
        std::cerr << "Kernel failed: " << cudaGetErrorString(err) << std::endl;
        return -1;
    }

    // 7. 拷贝结果回主机
    std::vector<unsigned int> h_histo_gpu(numFineBins);
    cudaMemcpy(h_histo_gpu.data(), d_histo, numFineBins * sizeof(unsigned int), cudaMemcpyDeviceToHost);

    // 8. CPU 计算参考结果
    std::vector<unsigned int> h_histo_cpu;
    cpuHistogram(h_vals, h_histo_cpu, numFineBins);

    // 9. 打印部分结果用于观察
    printHistogram(h_histo_gpu, "GPU Histogram");
    printHistogram(h_histo_cpu, "CPU Histogram (Reference)");

    // 10. 验证结果：GPU 与 CPU 是否一致
    bool passed = true;
    for (int i = 0; i < numFineBins; i++) {
        if (h_histo_gpu[i] != h_histo_cpu[i]) {
            std::cerr << "Mismatch at bin " << i << ": GPU=" << h_histo_gpu[i]
                      << ", CPU=" << h_histo_cpu[i] << std::endl;
            passed = false;
            break;
        }
    }

    // 11. 输出验证结果
    if (passed) {
        std::cout << "\nTest PASSED: GPU and CPU histograms match!\n";
    } else {
        std::cout << "\nTest FAILED: GPU and CPU histograms differ!\n";
    }

    std::cout << "GPU Time: " << time_gpu << " ms" << std::endl;

    // 12. 清理资源
    cudaFree(d_vals);
    cudaFree(d_histo);

    return 0;
}